import os
import json
file_path = os.path.abspath(__file__)
dir_path = os.path.dirname(file_path)


def data_cleaning(data_dir):
    """
    JD_chat.txt -> JD_chat.json
    按session_id将对话拆开，保存为json格式
    每个dict:
    {
        "session_id": 会话id,
        "user_id": 用户id,
        "content": 对话内容,
    }
    对话内容形式：
    user:\txxxxx\nwaiter:\txxxxx\n...
    """
    ret = []
    last_session_id = ""
    chat_log = {
        "session_id": "",
        "user_id": "",
        "content": "",
    }

    cnt = 0
    with open(os.path.join(dir_path, data_dir), "r", encoding="utf-8") as f:
        for line in f:
            if cnt:
                try:
                    line = line.strip("\n").strip("\t").split("\t")
                    session_id, user_id, waiter_send, is_transfer, is_repeat= line[: 5]
                    content = " ".join(line[6: ])
                except:
                    print(cnt)
                    break

                tmp = "waiter:\t" if int(waiter_send) else "user:\t"
                tmp += content + "\n"
                
                if session_id == last_session_id:
                    chat_log["content"] += tmp
                else:
                    ret.append(chat_log)
                    chat_log = {
                        "session_id": session_id,
                        "user_id": user_id,
                        "content": tmp,
                    }
                    last_session_id = session_id
                
            else:
                pass
            cnt += 1
    return cnt-1, len(ret), ret

if __name__ == "__main__":
    data_dir = "JD_chat.txt"
    n, chats, ret = data_cleaning(data_dir=data_dir)
    print(f"总数据量{n}，对话数{chats}")
    with open(os.path.join(dir_path, "JD_chat.json"), "w", encoding="utf-8") as f:
        json.dump(ret, f, ensure_ascii=False, indent=4)